import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from orbitP.layers.cross_encoder import Encoder
from orbitP.layers.cross_decoder import Decoder
from orbitP.layers.cross_attn import FullAttention, AttentionLayer, TwoStageAttentionLayer
from orbitP.layers.cross_embed import DSW_embedding
from math import ceil
from orbitP.script import config
from orbitP.layers.FeatureEncoder import FeatureEncoder

class crossformer(nn.Module):
    def __init__(self, args):
        super(crossformer, self).__init__()
        self.data_dim = args.head_dim
        self.n_heads = args.k
        self.in_len = config.training_length
        self.out_len = config.predicting_length
        self.seg_len = args.seg_len
        self.win_size = args.win_size
        self.e_layers = args.num_layers
        self.baseline = False

        self.device = torch.device(config.device)

        # The padding operation to handle invisible sgemnet length
        self.pad_in_len = ceil(1.0 * self.in_len / self.seg_len) * self.seg_len
        self.pad_out_len = ceil(1.0 * self.out_len / self.seg_len) * self.seg_len
        self.in_len_add = self.pad_in_len - self.in_len

        # Embedding
        self.enc_value_embedding = DSW_embedding(self.seg_len, args.d_model)
        self.enc_pos_embedding = nn.Parameter(torch.randn(1, self.data_dim, (self.pad_in_len // self.seg_len), args.d_model))
        self.pre_norm = nn.LayerNorm(args.d_model)

        # Encoder
        self.encoder = Encoder(self.e_layers, self.win_size, args.d_model, self.n_heads, args.d_ff, block_depth = 1, \
                                    dropout = args.dropout,in_seg_num = (self.pad_in_len // self.seg_len), factor = args.factor)
        
        # Decoder
        self.dec_pos_embedding = nn.Parameter(torch.randn(1, self.data_dim, (self.pad_out_len // self.seg_len), args.d_model))
        self.decoder = Decoder(self.seg_len, self.e_layers + 1, args.d_model, self.n_heads, args.d_ff, args.dropout, \
                                    out_seg_num = (self.pad_out_len // self.seg_len), factor = args.factor)

        self.outputLayer = nn.Linear(self.data_dim, config.outputSize, bias=True)
        self.FeatureEncoder = FeatureEncoder(args)

    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
        x_enc = x_enc.clone()
        x_mark_enc = x_mark_enc.clone()
        x_dec = x_dec[:,:,config.axis:-1].clone()
        x_mark_dec = x_mark_dec.clone()
        input = self.FeatureEncoder(x_enc, x_mark_enc, x_dec, x_mark_dec)

        if (self.baseline):
            base = input.mean(dim = 1, keepdim = True)
        else:
            base = 0
        batch_size = input.shape[0]
        if (self.in_len_add != 0):
            input = torch.cat((input[:, :1, :].expand(-1, self.in_len_add, -1), input), dim = 1)

        input = self.enc_value_embedding(input)
        input += self.enc_pos_embedding
        input = self.pre_norm(input)
        
        output = self.encoder(input)

        dec_in = repeat(self.dec_pos_embedding, 'b ts_d l d -> (repeat b) ts_d l d', repeat = batch_size)
        predict_y = self.decoder(dec_in, output)
        predict_y = self.outputLayer(predict_y)


        return base + predict_y[:, :self.out_len, :]